package edu.gatech.cc.liam.marl.decpomdp;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Set;

import edu.gatech.cc.liam.core.Globals;
import edu.gatech.cc.liam.core.Pair;
import edu.gatech.cc.liam.core.PairComparable;
import edu.gatech.cc.liam.core.rl.models.CartesianIterator;
import edu.gatech.cc.liam.core.rl.models.decpomdp.DecMDP;
import edu.gatech.cc.liam.core.rl.models.decpomdp.DecMDPPolicy;
import edu.gatech.cc.liam.core.rl.models.decpomdp.DecMDPSingleStagePolicy;
import edu.gatech.cc.liam.geometry.linear.FeasibleSet;
import edu.gatech.cc.liam.geometry.linear.Halfspace;
import edu.gatech.cc.liam.geometry.linear.NPoint;
import edu.gatech.cc.liam.geometry.linear.programming.BinaryProb;
import edu.gatech.cc.liam.geometry.linear.programming.LinearProb;
import edu.gatech.cc.liam.marl.decpomdp.Belief.DVO;

public class DMDPValueIteration {
//	private final static double MIN_VALUE_IMPROVEMENT_TO_SEEK = 0.001;
	private double minSeekImprovement;

	DMDPValueFunction previousVF;
	DMDPValueFunction currentVF;

	DecMDP baseModel;
	private BinaryProb probSolver;
	
	private int iterationNum;
	private double maxChangeInValue = 0.0;
	private double avgChangeInValue = -1.0;
	
	private Set<Belief> currentBs;
	
	public DMDPValueIteration(DecMDP baseModel, ArrayList<ValueVector> initialVVs) {
		this.baseModel = baseModel;
		for(ValueVector v : initialVVs)
			v.nextActions = baseModel.makeRandomPolicy();
		this.currentVF = new DMDPValueFunction(baseModel, initialVVs);
		this.previousVF = null;
		// define linear program over valid joint-policies
		FeasibleSet fs = baseModel.noNature.makePolicyFeasibleSet();
		probSolver = new BinaryProb(fs);
	}
	
//	
//	private DMDPValueIteration(DMDPValueIteration toClone) {
//		super();
//		this.previousVF = toClone.previousVF;
//		this.currentVF.valueFunction = toClone.currentVF.valueFunction;
//		this.baseModel = toClone.baseModel;
//		this.defaultBasis = toClone.defaultBasis;
//		this.probSolver = toClone.probSolver;
//		this.fs = toClone.fs;
//		this.lastBFS = toClone.lastBFS;
//	}
	
//	public BeliefValueFunction getFullValueFunctionBackup(Set<Belief> beliefs) {
//		BeliefValueFunction newVF = new BeliefValueFunction(baseModel,
//			    new ArrayList<PolicyVector>());
//		for(Belief targetB : beliefs) {
//			PolicyVector support = findUpdatedBeliefSupport(targetB);
//			if(!hasVector(support))
//				newVF.valueVectors.add(support);
//		}
//		return newVF;
//	}

	public void runVI(Belief startBelief,
					  int maxBeliefs, double epsilon, BeliefGenerator bGen) {
		assert this.baseModel.discount < 1.0 - Globals.TOLERANCE;
		minSeekImprovement = epsilon;
		iterationNum = 0;
		currentBs = new HashSet<Belief>();
		currentBs.add(startBelief);
		while(currentBs.size() < maxBeliefs * 10) {
			assert checkBeliefs(currentBs);
			double itChange = Double.MAX_VALUE;
			while(itChange > epsilon) {
				itChange = 0.0;
				iterationNum++;
				makeCurrentPrevious();
				System.out.println("\n*********** Iteration " + iterationNum + 
								" with " + currentBs.size() + " beliefs and " +
								previousVF.getNumSupports() + " supports. ********");
				// a round of perseus style backup (improving actions)
				perseusBackup();
				updateChangeInValues();
				itChange += maxChangeInValue;
				System.out.println(" numNotSearched = " + numNotSearched() + " " + this);
				// followed by propagating value improvements throughout the value-function
				// without changing actions
				makeCurrentPrevious();
				int numImprovment = propagateValues();
				updateChangeInValues();
				itChange += maxChangeInValue;
				int numBImproved = 0;
				for(Belief b : currentBs) {
					numBImproved += b.hasImproved ? 1 : 0;
				}
				System.out.println("Propigated " + numImprovment + " times for a total of " +
									numBImproved + " updated beliefs. " + this);
				System.out.println("---end-" + iterationNum + "---");
				System.out.println("Value at start: " + 
									currentVF.findValueAtBelief(startBelief) +
									"   Utility of best policy to discount " + epsilon + " steps: " +
									currentVF.simulateBestPolicy(startBelief, epsilon * 0.00001));
//				assert(currentVF.findValueAtBelief(startBelief) <=
//						currentVF.simulateBestPolicy(startBelief, epsilon * 0.01) + (epsilon * 0.01));
			}
			// add more beliefs
			int beforeSize = currentBs.size();
			bGen.addBeliefs(currentBs);
			System.out.println("#########\n#########" + (currentBs.size() - beforeSize) + " beliefs added.");
			if(beforeSize == currentBs.size()) break; // no beliefs added, so end value iteration
		}
	}

	private boolean checkBeliefs(Set<Belief> bs) {
		for(Belief b : bs) {
			b.validate(this.baseModel.validTypeList);
		}
		return true;
	}
	
	private void makeCurrentPrevious() {
		currentVF.updateAllBeliefs(currentBs);
		previousVF = currentVF;
		currentVF = new DMDPValueFunction(baseModel);
	}
	
	public DecMDPPolicy getEpsilonGreedyPolicy(final double epsilonGreed, final DecMDPPolicy explorePolicy) {
		class RandomActionGen implements DecMDPPolicy {
			// for each player and type select an action randomly
			public DecMDPSingleStagePolicy getActions(Belief publicB) {
				if(Globals.rand.nextDouble() <= epsilonGreed) {
					return currentVF.getActions(publicB);
				} else {
					return explorePolicy.getActions(publicB);
				}
			}
		}
		return new RandomActionGen();
	}
	
	private void updateChangeInValues() {
		// assumes beliefs have current value
		maxChangeInValue = Double.NEGATIVE_INFINITY;
		double valueSum = 0.0;
		for(Belief b : currentBs) {
			double change = currentVF.findValueAtBelief(b) - previousVF.findValueAtBelief(b);
			assert change > -Globals.TOLERANCE;
			maxChangeInValue = Math.max(maxChangeInValue, change);
			valueSum += change;
		}
		avgChangeInValue = valueSum / currentBs.size();
	}
	
	private void perseusBackup() {
		// assumes beliefs have current value
		LinkedList<Belief> beliefsToCheckSupport = new LinkedList<Belief>();
		int numLPs = 0;
		int numVsImproved = 0;
		int counter = 0;
		// keep searching for supports until at least one is found, or all successor vvs are searched
		HashSet<ValueVector> newVVs = new HashSet<ValueVector>();
		LinkedList<Belief> unsupportedBs = new LinkedList<Belief>(currentBs);
		// shuffle beliefs to randomize search order
		Collections.shuffle(unsupportedBs);
		// mark all beliefs unsupported
		for(Belief b : unsupportedBs) {
			b.hasImproved = false;
		}
		// find support (value of best one step look ahead policy) for a belief
		// mark all improved beliefs as supported
		// continue until no unsupported beliefs are left
		while(numVsImproved == 0 && !unsupportedBs.isEmpty()) {
			if(unsupportedBs.size() < (((10-counter) * currentBs.size()) / 10)) {
				System.out.print("|" + counter);
				counter++;
			}
			Belief targetBelief = unsupportedBs.remove();
			assert this.currentBs.contains(targetBelief);
			if(targetBelief.hasImproved) {
				// belief has been improved already, so leave it removed
				continue;
				// shouldn't happen
			}
			ValueVector support = exploreASupp(targetBelief);
			numLPs++;
			if(support == null) {
				// this belief's successor searches are exhausted (don't remove other supported Bs)
				newVVs.add(targetBelief.getSupport());
			} else if(support == targetBelief.getSupport()) {
				// a better action was not found, so do nothing for now
				// add it back so we can search the next support
				assert this.currentBs.contains(targetBelief);
				unsupportedBs.add(targetBelief);
			} else {
				numVsImproved++;
				targetBelief.hasImproved = true;
				newVVs.add(support);
				removeSupportedB(unsupportedBs, beliefsToCheckSupport, support);
			}
		}
		for(Belief b : beliefsToCheckSupport) {
			double newValue = currentVF.findValueAtBelief(b);
			if(newValue + Globals.TOLERANCE < b.getValue()) {
				newVVs.add(b.getSupport());
			}
		}
		for(Belief b : unsupportedBs) {
			newVVs.add(b.getSupport());
		}
		for(ValueVector vv : newVVs) {
			currentVF.addSupport(vv);
		}
		int numBImproved = 0;
		for(Belief b : currentBs) {
			numBImproved += b.hasImproved ? 1 : 0;
		}
		System.out.print("\nBackup improved " + numBImproved + " beliefs with " + numVsImproved + 
				" VVs using " + numLPs + " LPs ");
	}


	/**
	 * Using the huerisitc of which successor VV would give the best value without a change in action
	 * we compute the best action, and return the resulting new VV or null if the current support is better
	 */
	private ValueVector exploreASupp(Belief b) {
		assert(b.validate(baseModel.validTypeList));
		// assume previous supports is the more recent complete set of supports
		// assumes beliefs have updated values
		DVO succToCheck;
		do {
			succToCheck = b.removeBestSucc();
		} while(succToCheck != null && succToCheck.getSuccVV().hasSeen.contains(b));
		if(succToCheck == null) {
			return null;
		}
		ValueVector succ = succToCheck.getSuccVV();
		double[] obj = succ.makeObjective(b);
		double lpUB = getLPUpperBound(obj);
		double utility;
		if(lpUB >= b.getValue() + minSeekImprovement) {
			boolean success = probSolver.findBoolSol(obj);//, null);
			assert success;
			utility = probSolver.getOptimalValue();
		} else {
			utility = Double.NEGATIVE_INFINITY;
		}
		succ.hasSeen.add(b);
		ValueVector newV = null;
		if(utility > b.getValue() + minSeekImprovement) {
			DecMDPSingleStagePolicy action = new DecMDPSingleStagePolicy(baseModel, probSolver.getOptimalPoint());
			newV = new ValueVector(action, succ, b, previousVF);
			assert(Math.abs(newV.getValue(b) - utility) < Globals.TOLERANCE);
		} else {
			// no improvement vector found
			newV = b.getSupport();
			assert(Math.abs(newV.getValue(b) - b.getValue()) < Globals.TOLERANCE);
		}
		return newV;
	}
	
	/**
	 * @return the value of the best joint-action without information constraints
	 * (as if the agents could communicate). This is fast to compute and acts as an upper-bound
	 * which allows us to skip many LPs. 
	 */
	private double getLPUpperBound(double[] obj) {
		double val = 0.0;
		for(int jTi=0; jTi<baseModel.noNature.numJTypes; jTi++) {
			double bestAVal = Double.NEGATIVE_INFINITY;
			for(int jAi=0; jAi < baseModel.numJActions; jAi++) {
				int varIndex = baseModel.noNature.getVarIndex(jTi, jAi);
				bestAVal = Math.max(bestAVal, obj[varIndex]);
			}
			 val += bestAVal;
		}
		return val;
	}
	
	private void removeSupportedB(LinkedList<Belief> unsupportedB,
								  LinkedList<Belief> beliefsToCheckSupport, ValueVector support) {
		Iterator<Belief> it = unsupportedB.iterator();
		while (it.hasNext()) {
			Belief b = it.next();
			double supportImprovement = support.getValue(b) - b.getValue();
			if(supportImprovement + Globals.TOLERANCE >= 0.0){
				it.remove();
				b.hasImproved = true;
			} else if (supportImprovement + (maxChangeInValue * baseModel.discount) >= 0.0) {
//				supportedB.add(b);
//				beliefsToCheckSupport.add(b);
			}
		}
	}

	/**
	 * @return the total improvement over all beliefs
	 */
	private int propagateValues() {
		// assume previous supports is the more recent complete set of supports
		// assumes beliefs have updated values
		int numImproved = 0;
		boolean hasImproved = true; // minSeekImprovement
		while(hasImproved) {
			hasImproved = false;
			for(Belief b : currentBs) {
				if(b.getSuccImprovement() > minSeekImprovement) {
					ValueVector newVV = b.getSuppFromBestSucc(previousVF);
					assert b.updateIfBetter(newVV);
					hasImproved = true;
					b.hasImproved = true;
					numImproved++;
					for(Belief b2 : currentBs) {
						if(b2.updateIfBetter(newVV))
							b2.hasImproved = true;
						b2.replaceSuccIfBetter(newVV);
					}
				}
			}
		}
		for(Belief b : currentBs) {
			this.currentVF.addSupport(b.getSupport());
		}
		return numImproved;
	}

	private int numNotSearched() {
		int numNotSearched = 0;
		for(ValueVector v : this.currentVF) {
			int numVHasNotSeen = currentBs.size() - v.hasSeen.size();
			assert numVHasNotSeen >= 0;
			numNotSearched += numVHasNotSeen;
		}
		return numNotSearched;
	}
	
	@Override
	public String toString() {
		return 	" " + currentVF.getNumSupports() + " Value Vectors, with max change " + 
				Globals.decimal5.format(this.maxChangeInValue) +
				" and avg change " + Globals.decimal5.format(avgChangeInValue);
	}
	
}
